import logging
from typing import Type

from app.core.config import settings
from app.models import (
    MagneticSpaceGroupDocument,
    PointGroupDocument,
    SpaceGroupDocument,
    SpinSpaceGroupDocument,
)
from beanie import Document, init_beanie
from pymongo import AsyncMongoClient
from pymongo.asynchronous.database import AsyncDatabase

logger = logging.getLogger(__name__)


class MongoDatabase:
    """用于管理 MongoDB 连接和操作的类"""

    hsp_db: AsyncDatabase | None = None

    async def initialize(self):
        self.client = AsyncMongoClient(settings.MONGO_URI, uuidRepresentation="standard")
        self.hsp_db = self.client[settings.MONGO_DB_NAME]

        await init_beanie(
            database=self.hsp_db,
            document_models=[
                PointGroupDocument,
                SpaceGroupDocument,
                MagneticSpaceGroupDocument,
                SpinSpaceGroupDocument,
            ],
        )

        logger.info(f"✅ 数据库已连接: {settings.MONGO_DB_NAME}")
        logger.info(f"📊 已注册 4 个 Document 模型")

    @staticmethod
    async def get_fast_count(model: Type[Document], query: dict | None = None) -> int:
        """
        为一个 Beanie 模型获取快速的文档计数

        - 如果没有提供查询条件，则使用快速但可能不精确的 `estimated_document_count`
        - 如果提供了查询条件，则使用精确的 `count_documents`

        Args:
            model: Beanie Document 类
            query: MongoDB 查询字典

        Returns:
            文档的数量
        """
        collection = model.get_pymongo_collection()

        if query:
            return await collection.count_documents(query)

        try:
            # motor 的方法是异步的，需要 await
            return await collection.estimated_document_count()
        except Exception:
            # 如果估算失败，回退到对全集进行精确计数
            return await collection.count_documents({})

    @staticmethod
    async def get_db_data_size(db: AsyncDatabase | None, timeout_ms: int = 5000) -> int:
        """
        异步获取数据库统计信息，带超时处理

        Args:
            db: AsyncIOMotorDatabase 实例
            timeout_ms: 超时时间（毫秒），默认 5000ms

        Returns:
            数据库数据大小（字节），失败时返回 0
        """
        if db is None:
            return 0

        try:
            dbStats = await db.command("dbStats", maxTimeMS=timeout_ms)
            return dbStats.get("dataSize", 0)
        except Exception as e:
            logger.warning(f"获取数据库统计信息失败: {e}")
            return 0


# 创建数据库管理实例（单例）
mongodb = MongoDatabase()
